import json
import pickle

import numpy as np
import torch

from ImageUtils import list_to_tensor


class Knn_for_single:
    def __init__(self, dataset):
        self.action_list = []
        self.load_dataset(dataset)
        self.load_json_data()

    def load_json_data(self):
        # let action_list become as [[id, class, flatten_poses]]

        # read pose_list from json_address, and change pose_list to poses_flatten
        for action in self.action_list:
            address = action[2]
            with open(address, 'r') as f:
                data = json.load(f)

            pose_list = []

            for frame in data.values():
                frame[0]['keypoints'] = list_to_tensor(frame[0]['keypoints'])
                pose_list.append(frame[0]['keypoints'].reshape(1, 17, 2).numpy())

            # cat pose_list
            cat_pose = None
            for pose in pose_list:
                if cat_pose is None:
                    cat_pose = pose
                    continue
                else:
                    cat_pose = np.concatenate((cat_pose, pose))

            flatten_poses = cat_pose.flatten()
            action[2] = flatten_poses

    def load_dataset(self, dataset):
        id_list = dataset[0]
        class_list = dataset[1]
        json_address_list = dataset[2]
        for i in range(len(id_list)):
            self.action_list.append([id_list[i], class_list[i], json_address_list[i]])

    def predict(self, frame_list):
        # input: the list of frame, the format likes [frame0, frame1, ..., frame29]

        # first: change frame_list to pose_list, note: our frame only have a single pose.
        pose_list = []

        for frame in frame_list:
            human_pose = frame[0]['keypoints']  # human_pose: tensor,shape = (17,2)
            human_pose = np.array(human_pose).reshape(1, 17, 2)
            pose_list.append(human_pose)

        # cat pose_list and flatten
        cat_pose = None
        for pose in pose_list:
            if cat_pose is None:
                cat_pose = pose
                continue
            else:
                cat_pose = np.concatenate((cat_pose, pose))
        flatten_poses = cat_pose.flatten()

        # compute distance_list
        distance_list = []
        for action in self.action_list:
            dataset_poses = action[2]
            new_flatten_poses, dataset_poses = self.pad(flatten_poses, dataset_poses)
            distance = np.linalg.norm(new_flatten_poses - dataset_poses)
            distance_list.append([action[0], action[1], distance])

        neighborhood_list = sorted(distance_list, key=lambda x: x[2])
        return neighborhood_list

    def pad(self, row_data, target_data):
        # 此方法只用于flatten后的数据
        a = row_data.shape[0]
        b = target_data.shape[0]
        if a > b:
            row_data, target_data = target_data, row_data
            a, b = b, a
        change_row_data = np.pad(row_data, (0, b - a))
        return change_row_data, target_data

class Knn_for_mmaction2:
    def __init__(self) -> None:
        with open('mmaction2/result/203_predict.pkl', 'rb') as f:
            self.knn_list = pickle.load(f)
    
    def knn(self, feature):
        l = []
        for i, data in enumerate(self.knn_list):
            dis = np.linalg.norm(feature - data)
            l.append([i, dis])
        return sorted(l, key = lambda x:x[1])[:5]

if __name__ == '__main__':
    knn = Knn_for_mmaction2()
        